{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-12-01T12:45:10.158717Z",
     "start_time": "2024-12-01T12:45:10.153717Z"
    }
   },
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "from torchvision import transforms"
   ],
   "outputs": [],
   "execution_count": 38
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T12:45:10.876079Z",
     "start_time": "2024-12-01T12:45:10.693264Z"
    }
   },
   "cell_type": "code",
   "source": [
    "train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor())\n",
    "test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())"
   ],
   "id": "e8d1136124775c4d",
   "outputs": [],
   "execution_count": 39
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T12:45:11.289636Z",
     "start_time": "2024-12-01T12:45:11.278628Z"
    }
   },
   "cell_type": "code",
   "source": [
    "train_loader=torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)\n",
    "test_loader=torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=True)"
   ],
   "id": "11aa40485a7781eb",
   "outputs": [],
   "execution_count": 40
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T12:45:11.870235Z",
     "start_time": "2024-12-01T12:45:11.860238Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self,hidden_size=128,n_classes=10):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1 = nn.Linear(in_features=28*28, out_features=hidden_size)\n",
    "        self.fc2 = nn.Linear(in_features=hidden_size, out_features=n_classes)\n",
    "        self.softmax = nn.LogSoftmax(dim=1)\n",
    "        self.relu = nn.ReLU()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.relu(self.fc1(x))\n",
    "        x = self.relu(self.fc2(x))\n",
    "        x = self.softmax(x)\n",
    "        return x"
   ],
   "id": "f5aca513734ed5ab",
   "outputs": [],
   "execution_count": 41
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T13:19:11.887830Z",
     "start_time": "2024-12-01T13:19:11.877839Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model = Net()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "optimizerM = torch.optim.SGD(model.parameters(), lr=0.001,momentum=0.1)"
   ],
   "id": "9f372d558be84795",
   "outputs": [],
   "execution_count": 85
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T13:05:47.881221Z",
     "start_time": "2024-12-01T13:05:47.873521Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def Gradmul(grads,weight):\n",
    "    for grad in grads:\n",
    "        grad.mul_(weight)\n",
    "    #print(grads.shape)\n",
    "    return grads\n",
    "\n",
    "def Gradadd(gradsL,gradsR):\n",
    "    for gradL, gradR in zip(gradsL, gradsR):\n",
    "        #print(gradL.shape,gradR.shape)\n",
    "        gradL.add_(gradR)\n",
    "    return gradsL"
   ],
   "id": "11b497fb228e0ce5",
   "outputs": [],
   "execution_count": 77
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T13:07:03.366775Z",
     "start_time": "2024-12-01T13:07:03.354727Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def trainH(model, train_loader, optimizer,epochs,movment=0.1):\n",
    "    model.train()\n",
    "    for epoch in range(epochs):\n",
    "        flag=1\n",
    "        for x, y in train_loader:\n",
    "            x=x.view(x.shape[0], -1) \n",
    "            output = model(x)\n",
    "            loss = criterion(output, y)\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            for pram in model.parameters():     #rmsprop原理差不多，但时间来不及了:)\n",
    "                if flag==1:\n",
    "                    mem=pram.grad\n",
    "                    flag=0\n",
    "                else:\n",
    "                    #print(pram.grad.shape)\n",
    "                    pram.grad=Gradadd(pram.grad,Gradmul(mem,movment))\n",
    "                    mem=Gradadd(Gradmul(pram.grad,0.5),Gradmul(mem,0.5))\n",
    "                break\n",
    "            optimizer.step()\n",
    "        print('Epoch: {}, Loss: {}'.format(epoch, loss.item()))"
   ],
   "id": "5df957850db586e1",
   "outputs": [],
   "execution_count": 80
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T13:19:17.474926Z",
     "start_time": "2024-12-01T13:19:17.468508Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def train(model, train_loader, optimizer,epochs):\n",
    "    model.train()\n",
    "    for epoch in range(epochs):\n",
    "        for x, y in train_loader:\n",
    "            x=x.view(x.shape[0], -1) \n",
    "            output = model(x)\n",
    "            loss = criterion(output, y)\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        print('Epoch: {}, Loss: {}'.format(epoch, loss.item()))"
   ],
   "id": "55371029090c1fdc",
   "outputs": [],
   "execution_count": 86
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T13:10:09.252001Z",
     "start_time": "2024-12-01T13:07:04.042129Z"
    }
   },
   "cell_type": "code",
   "source": "trainH(model,train_loader,optimizer,10)",
   "id": "88dea49077f034c5",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, Loss: 0.8937862515449524\n",
      "Epoch: 1, Loss: 0.8034133315086365\n",
      "Epoch: 2, Loss: 0.5451378226280212\n",
      "Epoch: 3, Loss: 0.7688862681388855\n",
      "Epoch: 4, Loss: 1.0160408020019531\n",
      "Epoch: 5, Loss: 0.5967958569526672\n",
      "Epoch: 6, Loss: 0.6998957395553589\n",
      "Epoch: 7, Loss: 0.5075684189796448\n",
      "Epoch: 8, Loss: 0.7951298356056213\n",
      "Epoch: 9, Loss: 0.3651370108127594\n"
     ]
    }
   ],
   "execution_count": 81
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T13:21:14.846830Z",
     "start_time": "2024-12-01T13:19:20.276804Z"
    }
   },
   "cell_type": "code",
   "source": "train(model,train_loader,optimizerM,10) #不是，哥们",
   "id": "748421a663eaa67c",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, Loss: 2.2424185276031494\n",
      "Epoch: 1, Loss: 2.120954751968384\n",
      "Epoch: 2, Loss: 2.005002021789551\n",
      "Epoch: 3, Loss: 1.5737062692642212\n",
      "Epoch: 4, Loss: 1.7346421480178833\n",
      "Epoch: 5, Loss: 1.3863043785095215\n",
      "Epoch: 6, Loss: 1.3257347345352173\n",
      "Epoch: 7, Loss: 1.492504596710205\n",
      "Epoch: 8, Loss: 1.5455724000930786\n",
      "Epoch: 9, Loss: 1.082837700843811\n"
     ]
    }
   ],
   "execution_count": 87
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "2b84989d4edf57df"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
